
#
# INPUT: ['NaCl2', 'H2(NO)2', 'H2(NaO)3', 'H2(NaO2)3', 'H2(ClN)4(NO)3']
#

#
# OUTPUT:
# NaCl2 - NaClCl - 93.89517
# H2(NO)2 - HHNONO - 62.02808
# H2(NaO)3 - HHNaONaONaO - 118.98339
# H2(NaO2)3 - HHNaOONaOONaOO - 166.98159
# H2(ClN)4(NO)3 - HHClNClNClNClNNONONO - 289.87178
#

import re, itertools
from pyparsing import nestedExpr
elements = {
    'Ac': 227.000000,
    'Ag': 107.868000,
    'Al': 26.981540,
    'Am': 243.000000,
    'Ar': 39.948000,
    'As': 74.921600,
    'At': 210.000000,
    'Au': 196.966500,
    'B': 10.811000,
    'Ba': 137.330000,
    'Be': 9.012180,
    'Bi': 208.980400,
    'Br': 79.904000,
    'C': 12.011000,
    'Ca': 40.078000,
    'Cd': 112.410000,
    'Ce': 140.120000,
    'Cl': 35.452700,
    'Co': 58.933200,
    'Cr': 51.996000,
    'Cs': 132.905400,
    'Cu': 63.546000,
    'Dy': 162.500000,
    'Er': 167.260000,
    'Eu': 151.965000,
    'F': 18.998400,
    'Fe': 55.847000,
    'Fr': 223.000000,
    'Ga': 69.723000,
    'Gd': 157.250000,
    'Ge': 72.610000,
    'H': 1.007940,
    'He': 4.002600,
    'Hf': 178.490000,
    'Hg': 200.590000,
    'Ho': 164.930300,
    'I': 126.904500,
    'In': 114.820000,
    'Ir': 192.220000,
    'K': 39.098300,
    'Kr': 83.800000,
    'La': 138.905500,
    'Li': 6.941000,
    'Lu': 174.967000,
    'Mg': 24.305000,
    'Mn': 54.938000,
    'Mo': 95.940000,
    'N': 14.006700,
    'Na': 22.989770,
    'Nb': 92.906400,
    'Nd': 144.240000,
    'Ne': 20.179700,
    'Ni': 58.693400,
    'Np': 237.048200,
    'O': 15.999400,
    'Os': 190.200000,
    'P': 30.973760,
    'Pa': 231.035900,
    'Pb': 207.200000,
    'Pd': 106.420000,
    'Pm': 145.000000,
    'Po': 209.000000,
    'Pr': 140.907700,
    'Pt': 195.080000,
    'Pu': 244.000000,
    'Ra': 226.025400,
    'Rb': 85.467800,
    'Re': 186.207000,
    'Rh': 102.905500,
    'Rn': 222.000000,
    'Ru': 101.070000,
    'S': 32.066000,
    'Sb': 121.757000,
    'Sc': 44.955900,
    'Se': 78.960000,
    'Si': 28.085500,
    'Sm': 150.360000,
    'Sn': 118.710000,
    'Sr': 87.620000,
    'Ta': 180.947900,
    'Tb': 158.925300,
    'Tc': 98.000000,
    'Te': 127.600000,
    'Th': 232.038100,
    'Ti': 47.880000,
    'Tl': 204.383000,
    'Tm': 168.934200,
    'U': 238.029000,
    'V': 50.941500,
    'W': 183.850000,
    'Xe': 131.290000,
    'Y': 88.905900,
    'Yb': 173.040000,
    'Zn': 65.390000,
    'Zr': 91.224000,
}

element_strs = sorted(elements.keys(), key=len, reverse=True)

def expand(atom):
    parsed = nestedExpr('(',')').parseString(atom).asList()[0]
    
    expanded_list = []
    for i in range(len(parsed)):
        if i < len(parsed)-1 and isinstance(parsed[i+1], str) and parsed[i+1].isdigit():
            expanded_list += parsed[i]*int(parsed[i+1])
        elif not parsed[i].isdigit():
            expanded_list.append(parsed[i])
    
    expanded_formula = ''
    for expr in expanded_list:
        tokens = [t for t in re.split(r'(%s)' % r'|'.join(elements),expr) if t]
        expanded_subformula=''
        for i in range(len(tokens)):
            if i < len(tokens)-1 and tokens[i+1].isdigit():
                expanded_subformula += tokens[i]*int(tokens[i+1])
            elif not tokens[i].isdigit():
                expanded_subformula += tokens[i]
        expanded_formula += expanded_subformula
    
    if any(str(n) in expanded_formula for n in range(0,10)):
        tokens = [t for t in re.split(r'(%s)' % r'|'.join(element_strs),expanded_formula) if t]
        expanded_formula=''
        for i in range(len(tokens)):
            if i < len(tokens)-1 and tokens[i+1].isdigit():
                expanded_formula += tokens[i]*int(tokens[i+1])
            elif not tokens[i].isdigit():
                expanded_formula += tokens[i]
        
    return expanded_formula

def mass(atom):
    tokens = [t for t in re.split(r'(%s)' % r'|'.join(element_strs),atom) if t]
    return sum(elements[t] for t in tokens)

atoms = ['NaCl2', 'H2(NO)2', 'H2(NaO)3', 'H2(NaO2)3', 'H2(ClN)4(NO)3']
for a in atoms:
    expanded = expand('(%s)' % a)
    print a + ' - ' + expanded + ' - ' + str(mass(expanded))
